#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <mpi.h>

// find the grid cells covered by particle at x & compute the CIC weights
void CIC_weights(float x[], int ng, float *xa, int ind[3][2], float w[3][2])
{
  int i;
  float dx = xa[1] - xa[0];

  for (i=0; i<3; i++) {
    // first make sure that positions fall in box volume
    while (x[i] < xa[0]-dx/2) x[i] += ng*dx;
    while (x[i] > xa[ng-1]+dx/2) x[i] -= ng*dx;

    // now figure out nearest grid cells
    ind[i][0] = (int)floor((x[i]-xa[0])/dx);
    if (ind[i][0]<0 || ind[i][0]>=ng-1) {
      ind[i][0]=ng-1; ind[i][1]=0;
      if (x[i]<xa[ind[i][0]]) 
	w[i][1] = 1 - (w[i][0] = (xa[ind[i][1]]-x[i])/dx); 
      else w[i][0] = 1 - (w[i][1]=(x[i]-xa[ind[i][0]])/dx);
    } else {
      ind[i][1]=ind[i][0]+1; 
      w[i][0] = 1 - (w[i][1]=(x[i]-xa[ind[i][0]])/dx);
    }
  }

  return;
}

FILE *fopen_check(char *filename, char *mode) 
{
  FILE *fp=fopen(filename,mode);
  if (fp==NULL) {
    fprintf(stderr,"Error opening file %s with mode %s\n",filename,mode);
    perror("fopen");
    MPI_Abort(MPI_COMM_WORLD, 1234); exit(EXIT_FAILURE);
  }
  return fp;
}

float ***tensor3(int n1, int n2, int n3)
{
  int i, j;
  float ***m;
  m=(float ***)malloc(n1*sizeof(float **));
  *m=(float **)malloc(n1*n2*sizeof(float *));
  if (!(**m=(float *)calloc(n1*n2*n3,sizeof(float)))) {
    fprintf(stderr, "Error: unable to allocate %d*%d*%d tensor\n",n1,n2,n3);
    MPI_Abort(MPI_COMM_WORLD, 1234); exit(EXIT_FAILURE);
  }
  for(i=0;i<n1;i++) {
    if (i) {m[i]=m[i-1]+n2; *(m[i]) = *(m[i-1]) + n2*n3;}
    for (j=1; j<n2; j++) m[i][j] = m[i][j-1]+n3;
  }
  return m;
}

int main(int argc, char *argv[])
{
  float ***g, *gsum=NULL, *xa, x[3], d;
  long n3=0, npart, ng3;
  float w[3][2];
  int f, p, i, j, k, ind[3][2], nfile=128, ng=128;
  int rank, nproc;
  FILE *fp;
  char filename[200];
  double boxsize;

  MPI_Init(&argc, &argv);
  MPI_Comm_size(MPI_COMM_WORLD, &nproc);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);

  if (argc < 2 || argc > 4) {
    if (rank==0) 
      fprintf(stderr,"Usage: %s particlefile [outfile] np1\n",argv[0]);
    MPI_Finalize();
    exit(0);
  }

  nfile = nproc; // follow this convention for the IC/gadget/bindelta codes all the way!
  f = rank;

  if (nfile==1) sprintf(filename, "%s", argv[1]);
  else sprintf(filename, "%s.%d", argv[1],f); // previous version had .%02d but for gadget this is the filname naming convention.
  fp=fopen_check(filename, "rb");
 // get p-- the # particles in the current filename  
  fseek(fp, 8, SEEK_SET);
  fread(&p, 4, 1, fp); // this gives the total # particles in this file, checked!
//  printf("particles: %d in file %d \n", p,f);
  npart = p;
  //ng = rint(atoi(argv[3])/8);
  ng = rint(atoi(argv[3]));
  fseek(fp, 132, SEEK_SET);
  fread(&boxsize, 8, 1, fp); //gives the correct boxsize---checked.
  boxsize = (float) boxsize; 

  MPI_Allreduce(&npart, &n3, 1, MPI_LONG, MPI_SUM, MPI_COMM_WORLD);
  
  xa = calloc(ng,4);
  for (i=0; i<ng; i++) xa[i] = (i+0.5)/ng;
  // create & zero the density grid
  g=tensor3(ng,ng,ng);
  ng3 = ng*ng*ng;
  d = ng3*1.0/n3;

  fseek(fp, 268, SEEK_SET);
  while (p>0) {
    --p;
    if (fread(x,4,3,fp) != 3) {
      fprintf(stderr, "Error reading particle from file %s\n", filename);
      perror("fread"); exit(EXIT_FAILURE);
    }
    x[0]/=boxsize;x[1]/=boxsize;x[2]/=boxsize;
    CIC_weights(x,ng,xa,ind,w);
    // now add densities to cells with appropriate weights
    for (i=0; i<2; i++) for (j=0; j<2; j++) for (k=0; k<2; k++) 
	g[ind[0][i]][ind[1][j]][ind[2][k]] += d * w[0][i] * w[1][j] * w[2][k];
  }
  fclose(fp);
  free(xa);

  if (rank==0) gsum = calloc(ng3, sizeof(float));
  MPI_Reduce(**g, gsum, ng3, MPI_FLOAT, MPI_SUM, 0, MPI_COMM_WORLD);
  free(**g); free(*g); free(g);

  MPI_Finalize();

  if (rank == 0) {
    for (i=0; i<ng3; i++) gsum[i] -= 1;
    if (argc==4) fp = fopen_check(argv[argc-2],"wb");
    else fp = stdout;
    fwrite(gsum, sizeof(float), ng3, fp);
    //also write a text for for checking the numbers
    if (argc==4) fclose(fp);
  }

  return 0;
}
